#!/usr/bin/env python

import os
import struct
import sys
import re
from optparse import OptionParser
import subprocess
from shutil import move, copy, rmtree
import filecmp

def_path="klee/klee-last/"
file_tmp_dir="tmp_input/"
afl_dir="fuzz/"
afl_in_dir="in"
afl_crash_dir="crashing"
afl_config=afl_dir+"run_afl.sh"

version_no=3
class KTestError(Exception):
	 pass

class KTest:
	@staticmethod
	def fromfile(path):
		if not os.path.exists(path):
			print("ERROR: file %s not found" % (path))
			sys.exit(1)
			
		f = open(path,'rb')
		hdr = f.read(5)
		if len(hdr)!=5 or (hdr!=b'KTEST' and hdr != b"BOUT\n"):
			raise KTestError('unrecognized file')
		version, = struct.unpack('>i', f.read(4))
		if version > version_no:
			raise KTestError('unrecognized version')
		numArgs, = struct.unpack('>i', f.read(4))
		args = []
		for i in range(numArgs):
			size, = struct.unpack('>i', f.read(4))
			args.append(str(f.read(size).decode(encoding='ascii')))
			
		if version >= 2:
			symArgvs, = struct.unpack('>i', f.read(4))
			symArgvLen, = struct.unpack('>i', f.read(4))
		else:
			symArgvs = 0
			symArgvLen = 0

		numObjects, = struct.unpack('>i', f.read(4))
		objects = []
		for i in range(numObjects):
			size, = struct.unpack('>i', f.read(4))
			name = f.read(size)
			size, = struct.unpack('>i', f.read(4))
			bytes = f.read(size)
			objects.append( (name,bytes) )

		# Create an instance
		b = KTest(version, args, symArgvs, symArgvLen, objects)
		# Augment with extra filename field
		b.filename = path
		return b

	def __init__(self, version, args, symArgvs, symArgvLen, objects):
		self.version = version
		self.symArgvs = symArgvs
		self.symArgvLen = symArgvLen
		self.args = args
		self.objects = objects

		# add a field that represents the name of the program used to
		# generate this .ktest file:
		program_full_path = self.args[0]
		program_name = os.path.basename(program_full_path)
		# sometimes program names end in .bc, so strip them
		if program_name.endswith('.bc'):
		 program_name = program_name[:-3]
		self.programName = program_name


class Command:
	def __init__(self, arg0):
		self.binary = arg0
		self.args = []
		self.files = []
		self.crashing = False
		self.input_dir = afl_in_dir
		self.input_file_size = 0
		self.master = False

	def __str__(self):
		ret_val = self.binary+" "
		for a in self.args:
			ret_val += a+" "
		for f in self.files:
			ret_val += f+" "
		return ret_val

	def __repr__(self):
		ret_val = self.binary+" "
		for a in self.args:
			ret_val += str(a)+" "
		for f in self.files:
			ret_val += str(f)+" "
		return ret_val

	def callable(self):
		call = []
		call.append(self.binary)
		for a in self.args:
			call.append(str(a))
		for f in self.files:
			call.append(file_tmp_dir+str(f))
		return call

	fuzzer_id = 1

	def afl_call(self):
		if(self.crashing):
			call = "afl-fuzz -C -i "+afl_crash_dir+" -o outputs "						
		else:
			call = "afl-fuzz -i in -o out "

		if (self.master):
			call += "-M fuzzer0 "
		else:
			call += "-S fuzzer"+str(self.fuzzer_id)+" "
			Command.fuzzer_id += 1

		call += self.binary+" "
		for a in self.args:
			call += str(a)+" "
		call += "@@ "*self.input_file_size 
		#call += str(self.files)
		return call

	def update_files(self, rep_file, new_file):
		for index, used_file in enumerate(self.files):
			if (used_file == rep_file):
				#print "replace: "+used_file+" -> "+new_file
				self.files[index] = new_file

def testingBinary(cmd):
	ret_value = -1

	try:
		process = subprocess.Popen(cmd.callable(), shell=False, 
			stdout=open(os.devnull, 'wb'), stderr=open(os.devnull, 'wb'))
		process.wait()
		ret_value = process.returncode
	except:
		#print cmd.callable()
		print "ERROR on subprocess! (testing binary)"
		print cmd.callable()

	return ret_value


	 
def main(args):
	op = OptionParser("usage: %prog files")
	# op.add_option('','--write-ints', dest='writeInts', action='store_true',
	#               default=False,
	#               help='convert 4-byte sequences to integers')
	opts,args = op.parse_args()
	
	ktest_filelist = args[1:]

	# if not dir given, try default path: klee/klee-last/*.ktest
	if not args:
		for i in os.listdir(def_path):
			if i.endswith(".ktest"): 
				ktest_filelist.append(def_path+i)
    

	if len(ktest_filelist) < 1:
		print "no .ktest files found"
		sys.exit()
	
	cmd_list = []

	if not os.path.exists(file_tmp_dir):
		os.makedirs(file_tmp_dir)

	for file in ktest_filelist:
		b = KTest.fromfile(file)
		
		binary = b.args[0]
		if binary[-3:] == ".bc": binary = binary[:-3]	# remove .bc extension
		cmd = Command(binary)							# save binary 

		input_count = 0

		for i,(name,data) in enumerate(b.objects):
			# parsing data to afl-inputs
			#print name

			# filter version number and other garbage
			if name in ["model_version", "stdin" ]:
				continue				
			if "-stat" in name:
				continue

			#filename = "fuzz/klee_inputs/"+file.split("/",1)[1][:-6]

			if name in ["n_args"]:
				argc = struct.unpack('i',data)[0]
				#print str(argc) + " args call"
				continue
			if "arg" in name:
				#print name + " = " + data
				#print data.split("\x00")[0]
				cmd.args.append(data.split("\x00")[0])	# handle strings like C char* (null termianted)
				continue

			if "-data" in name:
				#print "inputs file with data from: "+os.path.basename(file)
				filename = os.path.basename(file)
				if(filename[-6:] == ".ktest"):		# remove .ktest file extension
					filename = filename[:-6]

				cmd.files.append(filename)
				cmd.input_file_size += 1
				#print "write to "+file_tmp_dir+filename
				f = open(file_tmp_dir+filename,'w')
				f.write(data)
				f.close()

				print data

				continue

			# if you come here, we have an unhandled case
			print "unhandled case detected for: "+name

		#print cmd
		cmd_list.append(cmd)

	for cmd in cmd_list:
		print cmd
		# for filename in cmd.files:
		# 	f = open(file_tmp_dir+filename,'r')
		# 	print filename + " = -" + f.read()+"-"


if __name__=='__main__':
	main(sys.argv)